import os
import sys
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image, ImageTk
import imagehash
import numpy as np
from collections import defaultdict
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from sklearn.metrics.pairwise import cosine_similarity
import threading
import queue

# 设置matplotlib支持中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

class ImageSearchEngine:
    def __init__(self):
        self.image_hashes = {}
        self.image_histograms = {}
        
    def extract_features(self, image_path):
        """提取图像特征：哈希值和颜色直方图"""
        try:
            with Image.open(image_path) as img:
                # 计算感知哈希值
                phash = imagehash.phash(img)
                
                # 计算颜色直方图
                hist = img.histogram()
                if len(hist) > 768:  # 如果是RGBA图像，只取RGB部分
                    hist = hist[:768]
                
                # 归一化直方图
                hist = np.array(hist, dtype=np.float32)
                hist = hist / np.sum(hist) if np.sum(hist) > 0 else hist
                
                return phash, hist
        except Exception as e:
            print(f"处理图片 {image_path} 时出错: {e}")
            return None, None
    
    def build_index(self, directory, progress_callback=None):
        """构建图片库索引"""
        print(f"正在扫描目录: {directory}")
        count = 0
        image_files = []
        
        # 先收集所有图片文件
        for root, _, files in os.walk(directory):
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')):
                    image_files.append(os.path.join(root, file))
        
        total = len(image_files)
        
        # 处理每个图片文件
        for i, image_path in enumerate(image_files):
            phash, hist = self.extract_features(image_path)
            
            if phash is not None and hist is not None:
                self.image_hashes[image_path] = phash
                self.image_histograms[image_path] = hist
                count += 1
                
            # 更新进度
            if progress_callback:
                progress_callback(i + 1, total)
        
        print(f"已索引 {count} 张图片")
        return count
    
    def find_similar_images(self, query_image_path, top_n=5, hash_weight=0.7, hist_weight=0.3):
        """查找相似图片，返回相似度最高的top_n张图片"""
        if not os.path.exists(query_image_path):
            print(f"查询图片不存在: {query_image_path}")
            return []
        
        # 提取查询图片的特征
        query_phash, query_hist = self.extract_features(query_image_path)
        
        if query_phash is None or query_hist is None:
            print("无法提取查询图片的特征")
            return []
        
        # 计算相似度
        similarities = []
        
        for image_path, stored_phash in self.image_hashes.items():
            # 跳过查询图片自身
            if os.path.abspath(image_path) == os.path.abspath(query_image_path):
                continue
            
            # 计算哈希相似度 (值越小越相似)
            hash_distance = query_phash - stored_phash
            hash_similarity = 1.0 - (hash_distance / 64.0)  # 将汉明距离转换为相似度
            
            # 计算直方图相似度
            hist_similarity = cosine_similarity(
                [self.image_histograms[image_path]], 
                [query_hist]
            )[0][0]
            
            # 综合加权相似度
            combined_similarity = hash_weight * hash_similarity + hist_weight * hist_similarity
            similarities.append((image_path, combined_similarity))
        
        # 按相似度排序
        similarities.sort(key=lambda x: x[1], reverse=True)
        
        return similarities[:top_n]

class ImageSearchGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("图片相似度搜索工具")
        self.root.geometry("1200x800")
        
        self.engine = ImageSearchEngine()
        self.gallery_path = ""
        self.query_image_path = ""
        self.results = []
        self.top_n = 5
        
        # 创建线程间通信的队列
        self.queue = queue.Queue()
        
        self.create_widgets()
    
    def create_widgets(self):
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 左侧控制面板
        control_frame = ttk.LabelFrame(main_frame, text="控制面板", padding="10")
        control_frame.pack(side=tk.TOP, fill=tk.X, padx=5, pady=5)
        
        # 图库路径选择
        ttk.Label(control_frame, text="图库路径:").pack(side=tk.LEFT, pady=5)
        self.gallery_var = tk.StringVar()
        gallery_entry = ttk.Entry(control_frame, textvariable=self.gallery_var, width=40)
        gallery_entry.pack(side=tk.LEFT, padx=5, pady=5)
        
        browse_btn = ttk.Button(control_frame, text="浏览...", command=self.browse_gallery)
        browse_btn.pack(side=tk.LEFT, pady=5)
        
        # 构建索引按钮
        self.index_btn = ttk.Button(control_frame, text="构建图库索引", command=self.build_index)
        self.index_btn.pack(side=tk.LEFT, pady=10)
        
        # 索引进度条
        self.index_progress = ttk.Progressbar(control_frame, orient=tk.HORIZONTAL, length=100, mode='determinate')
        self.index_progress.pack(side=tk.LEFT,fill=tk.X, pady=5)
        
        self.progress_label = ttk.Label(control_frame, text="")
        self.progress_label.pack(side=tk.LEFT, pady=2)
        
        search_frame = ttk.LabelFrame(main_frame,text="选择图片",padding="10")
        search_frame.pack(side=tk.TOP, fill=tk.X, padx=5, pady=5)
        
        # 查询图片选择
        ttk.Label(search_frame, text="查询图片:").pack(side=tk.LEFT, pady=5)
        self.query_var = tk.StringVar()
        query_entry = ttk.Entry(search_frame, textvariable=self.query_var, width=40)
        query_entry.pack(side=tk.LEFT, padx=5, pady=5)
        
        browse_query_btn = ttk.Button(search_frame, text="浏览...", command=self.browse_query)
        browse_query_btn.pack(side=tk.LEFT, pady=5)
        
        # 结果数量选择
        ttk.Label(search_frame, text="显示结果数量:").pack(side=tk.LEFT, pady=5)
        self.top_n_var = tk.IntVar(value=5)
        top_n_combo = ttk.Combobox(search_frame, textvariable=self.top_n_var, values=[3, 5, 7, 10, 15], width=5)
        top_n_combo.pack(side=tk.LEFT, pady=5)
        
        # 搜索按钮
        self.search_btn = ttk.Button(search_frame, text="搜索相似图片", command=self.search_similar, state=tk.DISABLED)
        self.search_btn.pack(side=tk.LEFT, pady=20)
        
        # 右侧显示区域
        display_frame = ttk.LabelFrame(main_frame, text="结果显示", padding="10")
        display_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 查询图片显示区域
        query_display_frame = ttk.LabelFrame(display_frame, text="查询图片", padding="5")
        query_display_frame.pack(fill=tk.X, padx=5, pady=5)
        
        self.query_image_label = ttk.Label(query_display_frame, text="未选择图片")
        self.query_image_label.pack(padx=5, pady=5)
        
        # 结果显示区域
        results_frame = ttk.LabelFrame(display_frame, text="相似图片结果", padding="5")
        results_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 创建画布和滚动条
        self.canvas = tk.Canvas(results_frame)
        scrollbar = ttk.Scrollbar(results_frame, orient="vertical", command=self.canvas.yview)
        self.scrollable_frame = ttk.Frame(self.canvas)
        
        self.scrollable_frame.bind(
            "<Configure>",
            lambda e: self.canvas.configure(
                scrollregion=self.canvas.bbox("all")
            )
        )
        
        self.canvas.create_window((0, 0), window=self.scrollable_frame, anchor="nw")
        self.canvas.configure(yscrollcommand=scrollbar.set)
        
        scrollbar.pack(side="right", fill="y")
        self.canvas.pack(side="left", fill="both", expand=True)
    
    def browse_gallery(self):
        directory = filedialog.askdirectory(title="选择图库目录")
        if directory:
            self.gallery_path = directory
            self.gallery_var.set(directory)
    
    def browse_query(self):
        file_path = filedialog.askopenfilename(
            title="选择查询图片",
            filetypes=[("图片文件", "*.png;*.jpg;*.jpeg;*.bmp;*.webp")]
        )
        if file_path:
            self.query_image_path = file_path
            self.query_var.set(file_path)
            self.display_query_image()
            # 启用搜索按钮
            if self.engine.image_hashes:
                self.search_btn.config(state=tk.NORMAL)
    
    def display_query_image(self):
        try:
            # 加载并显示查询图片
            img = Image.open(self.query_image_path)
            img.thumbnail((400, 300))
            photo = ImageTk.PhotoImage(img)
            
            self.query_image_label.config(image=photo)
            self.query_image_label.image = photo
        except Exception as e:
            messagebox.showerror("错误", f"无法显示图片: {e}")
    
    def build_index(self):
        if not self.gallery_path:
            messagebox.showwarning("警告", "请先选择图库目录")
            return
        
        # 禁用按钮
        self.index_btn.config(state=tk.DISABLED)
        self.search_btn.config(state=tk.DISABLED)
        
        # 重置进度条
        self.index_progress["value"] = 0
        self.progress_label.config(text="")
        
        # 在新线程中构建索引
        threading.Thread(target=self._build_index_thread, daemon=True).start()
        
        # 开始轮询队列
        self.root.after(100, self._process_queue)
    
    def _build_index_thread(self):
        try:
            count = self.engine.build_index(self.gallery_path, self._update_progress)
            self.queue.put(("success", f"成功索引 {count} 张图片"))
            
            # 如果已选择查询图片，启用搜索按钮
            if self.query_image_path:
                self.queue.put(("enable_search", None))
        except Exception as e:
            self.queue.put(("error", f"构建索引时出错: {e}"))
    
    def _update_progress(self, current, total):
        self.queue.put(("progress", (current, total)))
    
    def _process_queue(self):
        while not self.queue.empty():
            try:
                msg = self.queue.get()
                if msg[0] == "progress":
                    current, total = msg[1]
                    self.index_progress["value"] = (current / total) * 100
                    self.progress_label.config(text=f"已处理: {current}/{total}")
                elif msg[0] == "success":
                    messagebox.showinfo("成功", msg[1])
                    self.index_btn.config(state=tk.NORMAL)
                elif msg[0] == "error":
                    messagebox.showerror("错误", msg[1])
                    self.index_btn.config(state=tk.NORMAL)
                elif msg[0] == "enable_search":
                    self.search_btn.config(state=tk.NORMAL)
            except queue.Empty:
                pass
        
        if self.index_btn["state"] == tk.DISABLED:
            self.root.after(100, self._process_queue)
    
    def search_similar(self):
        if not self.query_image_path:
            messagebox.showwarning("警告", "请先选择查询图片")
            return
        
        if not self.engine.image_hashes:
            messagebox.showwarning("警告", "请先构建图库索引")
            return
        
        # 禁用搜索按钮
        self.search_btn.config(state=tk.DISABLED)
        
        # 获取结果数量
        self.top_n = self.top_n_var.get()
        
        # 在新线程中搜索
        threading.Thread(target=self._search_thread, daemon=True).start()
        
        # 开始轮询队列
        self.root.after(100, self._process_search_queue)
    
    def _search_thread(self):
        try:
            self.results = self.engine.find_similar_images(self.query_image_path, self.top_n)
            print(self.results)
            self.queue.put(("search_success", None))
        except Exception as e:
            self.queue.put(("search_error", f"搜索时出错: {e}"))
    
    def _process_search_queue(self):
        while not self.queue.empty():
            try:
                msg = self.queue.get()
                if msg[0] == "search_success":
                    self.display_results()
                    self.search_btn.config(state=tk.NORMAL)
                elif msg[0] == "search_error":
                    messagebox.showerror("错误", msg[1])
                    self.search_btn.config(state=tk.NORMAL)
            except queue.Empty:
                pass
        
        if self.search_btn["state"] == tk.DISABLED:
            self.root.after(100, self._process_search_queue)
    
    def display_results(self):
        # 清空结果区域
        for widget in self.scrollable_frame.winfo_children():
            widget.destroy()
        
        if not self.results:
            ttk.Label(self.scrollable_frame, text="未找到相似图片").pack(padx=5, pady=5)
            return
        
        # 创建结果展示
        for i, (image_path, similarity) in enumerate(self.results, 1):
            result_frame = ttk.LabelFrame(self.scrollable_frame, text=f"相似度 #{i}: {similarity:.4f}", padding="5")
            result_frame.pack(fill=tk.X, padx=5, pady=5)
            
            # 左侧显示图片
            img_frame = ttk.Frame(result_frame)
            img_frame.pack(side=tk.LEFT, padx=5, pady=5)
            
            try:
                img = Image.open(image_path)
                img.thumbnail((200, 200))
                photo = ImageTk.PhotoImage(img)
                
                img_label = ttk.Label(img_frame, image=photo)
                img_label.image = photo
                img_label.pack()
            except Exception as e:
                ttk.Label(img_frame, text=f"无法显示图片: {e}").pack()
            
            # 右侧显示信息
            info_frame = ttk.Frame(result_frame)
            info_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=5, pady=5)
            
            ttk.Label(info_frame, text="图片路径:").pack(anchor=tk.W)
            path_label = ttk.Label(info_frame, text=image_path, wraplength=600)
            path_label.pack(anchor=tk.W, pady=2)
            
            # 打开图片按钮
            open_btn = ttk.Button(info_frame, text="打开图片", command=lambda path=image_path: self.open_image(path))
            open_btn.pack(anchor=tk.W, pady=5)
    
    def open_image(self, image_path):
        try:
            os.startfile(image_path)  # Windows系统
        except AttributeError:
            try:
                os.system(f'open "{image_path}"')  # macOS系统
            except:
                os.system(f'xdg-open "{image_path}"')  # Linux系统

def main():
    root = tk.Tk()
    app = ImageSearchGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()    